import os
import pydicom
import SimpleITK as sitk
from DicomAutoConvert.DicomDataCheck import DicomDataCheck


class ConvertErrorDefine:
    CHECK_ID_NOT_FOUND_ERROR = "无法搜索到检查号字段"
    STUDY_INFORMATION_NOT_FOUND_ERROR = "无法搜索到检查部位字段"
    STUDY_INFORMATION_NOT_MATCH_ERROR = "检查部位关键词匹配异常"
    FRAME_NUMBER_LOSS_ERROR = "帧缺失异常"
    NOT_CHEST_DATA_SOURCE_ERROR = "非胸部CT数据源异常(算法质量评估异常)"
    WINDOWS_LOSS_DATA_SOURCE_ERROR = "胸部窗缺失数据(或无法读取)异常"
    UNEXCEPTION_OCCUR_ERROR = "程序未定义错误: {}"


class RawDataCollection:
    """文件结构映射
        parent1:
            - dcm1
            - dcm2
        parent2:
            - dcm1
            - dcm2
    """
    def __init__(self, root_start: str):
        self.root_start = root_start
        self.root_dcm_dict = {}
        self.most_match_serial_id = {}

        self.BASIC_INFO_ACCESSION_NUMBER = (0x0008, 0x0050)
        self.BASIC_INFO_STUDY_DESCRIPTION = (0x0008, 0x1030)

    def save_error_to_csv(self, error_dcm_dict: dict, default_save_name: str = "error_dcm_files.csv"):
        with open(default_save_name, 'w', encoding='utf-8') as f:
            f.write("文件路径, 错误原因" + os.linesep)
            f.flush()

            for each_error in error_dcm_dict:
                f.write("{}, {}".format(each_error, error_dcm_dict[each_error]) + os.linesep)
            f.flush()
        return True

    def convert_clean_dcm_dict_to_nrrd(self, output_nrrd_path: str):
        os.makedirs(output_nrrd_path, exist_ok=True)

        for each_clean_key in self.root_dcm_dict:
            ds = pydicom.dcmread(os.path.join(each_clean_key, self.root_dcm_dict[each_clean_key][0]))
            check_id = ds.AccessionNumber
            serial_id = self.most_match_serial_id[each_clean_key]

            reader = sitk.ImageSeriesReader()
            series_file_names = reader.GetGDCMSeriesFileNames(each_clean_key, serial_id)
            reader.SetFileNames(series_file_names)
            dicom_series = reader.Execute()

            sitk.WriteImage(dicom_series, os.path.join(output_nrrd_path, check_id + ".nrrd"))
        return True

    def check_every_study(self):
        info_clean_dcm_dict = {}
        error_dcm_dict = {}

        # 第一部分过滤基础信息
        for each_key in self.root_dcm_dict:
            basic_error_info = self.check_basic_dcm_information(each_key, self.root_dcm_dict[each_key])
            if basic_error_info != "":
                error_dcm_dict[each_key] = basic_error_info
            else:
                info_clean_dcm_dict[each_key] = self.root_dcm_dict[each_key]

        # 第二部分过滤数据异常
        data_clean_dcm_dict = {}
        for each_key in info_clean_dcm_dict:
            basic_error_info = self.check_dcm_data_error(each_key, self.root_dcm_dict[each_key])
            if basic_error_info != "":
                error_dcm_dict[each_key] = basic_error_info
            else:
                data_clean_dcm_dict[each_key] = info_clean_dcm_dict[each_key]

        self.root_dcm_dict = data_clean_dcm_dict
        print("############### Cleaned DCM List ##############")
        print(self.root_dcm_dict.keys())
        return error_dcm_dict

    def check_dcm_data_error(self, root_dir: str, dcm_file_list: list):
        # TODO: 1. 判断是否能正常读取, 这也包含了窗是否丢失的判断?
        try:
            if not DicomDataCheck.sitk_windows_loss(root_dir):
                return ConvertErrorDefine.WINDOWS_LOSS_DATA_SOURCE_ERROR
        except Exception as e:
            return ConvertErrorDefine.UNEXCEPTION_OCCUR_ERROR.format(str(e))

        # TODO: 2. 判断帧丢失, 可能需要v1, v2组合
        try:
            if not DicomDataCheck.sitk_check_frame_loss_v2(root_dir, dcm_file_list):
                return ConvertErrorDefine.FRAME_NUMBER_LOSS_ERROR
        except Exception as e:
            return ConvertErrorDefine.UNEXCEPTION_OCCUR_ERROR.format(str(e))

        # TODO: 3. 判断帧内容是否合格
        try:
            quality_list, match_index = DicomDataCheck.sitk_data_useable_check(root_dir, dcm_file_list)
            if match_index == -1:
                return ConvertErrorDefine.NOT_CHEST_DATA_SOURCE_ERROR
            self.most_match_serial_id[root_dir] = quality_list[match_index][0]
        except Exception as e:
            return ConvertErrorDefine.UNEXCEPTION_OCCUR_ERROR.format(str(e))

        return ""

    def check_basic_dcm_information(self, root_dir: str, dcm_file_list: list):
        for each_dcm_file in dcm_file_list:
            ds = pydicom.dcmread(os.path.join(root_dir, each_dcm_file))
            if self.BASIC_INFO_ACCESSION_NUMBER not in ds:
                return ConvertErrorDefine.CHECK_ID_NOT_FOUND_ERROR
            if self.BASIC_INFO_STUDY_DESCRIPTION not in ds:
                return ConvertErrorDefine.STUDY_INFORMATION_NOT_FOUND_ERROR

            match_list_keys = ["Thorax", "chest", "CHEST", "胸部平扫", "胸部CT平扫"]
            for each_key in match_list_keys:
                if each_key in ds.StudyDescription:
                    return ""
            return ConvertErrorDefine.STUDY_INFORMATION_NOT_MATCH_ERROR

    def get_root_dcm_files(self):
        for root, dirs, files in os.walk(self.root_start):
            filter_dcm_file = []
            for each_file in files:
                if each_file.lower().endswith(".dcm"):
                    filter_dcm_file.append(each_file)

            if len(filter_dcm_file) == 0:
                continue

            # TODO: 文件夹可能存在多个脏数据
            self.root_dcm_dict[root] = filter_dcm_file


root_start = "/mnt/f1f2bfba-ca3d-406d-a60b-76f199e17d7f/imedFtpDir/newtest/DIYIPItest/ZONGGECHUANG"
collection = RawDataCollection(root_start)
collection.get_root_dcm_files()
error_dict_list = collection.check_every_study()
print(collection.save_error_to_csv(error_dict_list))
print(collection.convert_clean_dcm_dict_to_nrrd(os.path.join(os.path.dirname(os.path.abspath(__file__)), "outputnrrd")))
